# 提出问题--->文档加载--->文档分割--->embeddings词向量--->存储至向量数据库--->向量数据库根据问题检索数据--->生成含Prompt的问题--->将问题传递给LLM--->LLM给出答案

# 问答（Question Answering）

import os
from langchain.memory import ConversationBufferMemory
from langchain.vectorstores import Chroma
from langchain.embeddings.openai import OpenAIEmbeddings
from langchain.chat_models import ChatOpenAI
# 导入检索式问答链
from langchain.chains import RetrievalQA, ConversationalRetrievalChain
# 导入模板的检索式问答链
from langchain.prompts import PromptTemplate

api_key = "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"
os.environ['OPENAI_API_KEY'] = "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"

# 加载在之前已经进行持久化的向量数据库
persist_directory = 'docs/chroma/cs229_lectures/'
embedding = OpenAIEmbeddings()
vectordb = Chroma(persist_directory=persist_directory, embedding_function=embedding)
# 使用 ChatGPT3.5，温度设置为0
chat = ChatOpenAI(api_key=api_key, model_name="gpt-3.5-turbo", temperature=0)


# 普通检索式问答链
def qa_simple():
    # 声明一个检索式问答链
    qa_chain = RetrievalQA.from_chain_type(
        chat,
        retriever=vectordb.as_retriever()
    )
    # 可以以该方式进行检索问答
    question1 = "Is probability a class topic?"
    result1 = qa_chain({"query": question1})
    print(result1["result"])
    question2 = "why are those prerequesites needed?"
    result2 = qa_chain({"query": question2})
    print(result2["result"])

# 基于模板的检索式问答链 - 这种方法非常好，因为它只涉及对语言模型的一次调用。然而，它也有局限性，即如果文档太多，可能无法将它们全部适配到上下文窗口中
def qa_prompt_template():
    # Build prompt
    template = """Use the following pieces of context to answer the question at the end. If you don't know the answer, just say that you don't know, don't try to make up an answer. Use three sentences maximum. Keep the answer as concise as possible. Always say "thanks for asking!" at the end of the answer. 
    {context}
    Question: {question}
    Helpful Answer:"""
    QA_CHAIN_PROMPT = PromptTemplate.from_template(template)
    # Run chain
    qa_chain = RetrievalQA.from_chain_type(
        chat,
        retriever=vectordb.as_retriever(),
        return_source_documents=True,
        chain_type_kwargs={"prompt": QA_CHAIN_PROMPT}
    )
    # 可以以该方式进行检索问答
    question = "Is probability a class topic?"
    result = qa_chain({"query": question})
    print(result["result"])


# 基于MapReduce的检索式问答链 - 当我们将之前的问题通过这个链进行运行时，我们可以看到这种方法的两个问题。第一，速度要慢得多。第二，结果实际上更差。根据给定文档的这一部分，对这个问题并没有明确的答案。这可能是因为它是基于每个文档单独回答的。因此，如果信息分布在两个文档之间，它并没有在同一上下文中获取到所有的信息。
def qa_map_reduce():
    qa_chain_mr = RetrievalQA.from_chain_type(
        chat,
        retriever=vectordb.as_retriever(),
        chain_type="map_reduce"
    )
    # 可以以该方式进行检索问答
    question = "Is probability a class topic?"
    result = qa_chain_mr({"query": question})
    print(result["result"])


# 基于Refine的检索式问答链 - 我们可以类似地设置链式类型为Refine。这是一种新的链式类型。Refine 文档链类似于 MapReduce 链，对于每一个文档，会调用一次 LLM，但有所改进的是，我们每次发送给 LLM 的最终提示是一个序列，这个序列会将先前的响应与新数据结合在一起，并请求得到改进后的响应。因此，这是一种类似于 RNN 的概念，我们增强了上下文，从而解决信息分布在不同文档的问题。
def qa_refine():
    qa_chain_mr = RetrievalQA.from_chain_type(
        chat,
        retriever=vectordb.as_retriever(),
        chain_type="refine"
    )
    # 可以以该方式进行检索问答
    question = "Is probability a class topic?"
    result = qa_chain_mr({"query": question})
    print(result["result"])


def qa_memory_example():
    # 1、添加记忆功能
    memory = ConversationBufferMemory(
        memory_key="chat_history",  # 与 prompt 的输入变量保持一致。
        return_messages=True  # 将以消息列表的形式返回聊天记录，而不是单个字符串
    )
    # 2、声明一个检索式问答链
    retriever = vectordb.as_retriever()
    qa = ConversationalRetrievalChain.from_llm(
        chat,
        retriever=retriever,
        memory=memory
    )
    # 可以以该方式进行检索问答
    question1 = "Is probability a class topic?"
    result1 = qa({"question": question1})
    print(result1['answer'])
    question2 = "why are those prerequesites needed?"
    result2 = qa({"question": question2})
    print(result2['answer'])



if __name__ == '__main__':
    # qa_simple()
    # qa_prompt_template()
    # qa_map_reduce()
    # qa_refine()

    # 基本上，我们使用的链式（chain）没有任何状态的概念。它不记得之前的问题或之前的答案。为了实现这一点，我们需要引入内存，这是我们将在下一节中讨论的内容。
    qa_memory_example()

